package taskq

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/go-redis/redis/v8"
)

/**
* Author: dinglin
* Date: 2022/3/23 5:17 下午
 */
const (
	zSetKeyPrefix   = "tq_zset_"
	streamKeyPrefix = "tq_zset_stream_"
)

type ITaskQueue interface {
	AddTask(context.Context, *Task) error
	CancelTask(context.Context, string) error
	StartConsumer()
	Close()
}

type TaskQueueOption struct {
	Name       string
	zSetKey    string
	streamKey  string
	idAllocKey string
	Limiter    Limiter
	Interval   time.Duration
	r          redis.UniversalClient
	batchCount int
}

func NewTaskQueue(option *TaskQueueOption) ITaskQueue {
	if option.Name == "" {
		panic("task queue name missed")
	}
	queue := &TaskQueue{
		name:               option.Name,
		zSetKey:            option.zSetKey,
		streamKey:          option.streamKey,
		idAllocKey:         option.idAllocKey,
		limiter:            option.Limiter,
		interval:           option.Interval,
		r:                  option.r,
		batchCount:         option.batchCount,
		streamGroupName:    "tq",
		streamConsumerName: makeConsumerName(),
		exit:               make(chan bool),
		logger:             log.New(os.Stderr, "", log.LstdFlags|log.Lshortfile),
	}
	if queue.zSetKey == "" {
		queue.zSetKey = zSetKeyPrefix + option.Name
	}
	if queue.streamKey == "" {
		queue.streamKey = streamKeyPrefix + option.Name
	}
	if queue.limiter == nil {
		option.Limiter = &DefaultLimiter{}
	}
	if queue.interval == 0 {
		queue.interval = time.Millisecond * 100
	}
	if queue.batchCount == 0 {
		queue.batchCount = 10
	}
	queue.wg.Add(3)
	go queue.schedule()
	go queue.pending()
	go queue.cleanZombie()
	return queue
}

type TaskQueue struct {
	logger             *log.Logger
	name               string
	zSetKey            string
	streamKey          string
	idAllocKey         string
	streamGroupName    string
	streamConsumerName string
	limiter            Limiter
	interval           time.Duration
	r                  redis.UniversalClient
	batchCount         int
	exit               chan bool
	wg                 sync.WaitGroup
}

func (queue *TaskQueue) AddTask(ctx context.Context, task *Task) error {
	h := GetTask(task.Name)
	if h == nil {
		return fmt.Errorf("task(%s) handler not found", task.Name)
	}
	if task.Delay <= 0 {
		return queue.directlyAddTask(ctx, task)
	}
	absTime := time.Now().UTC().Add(task.Delay)
	pipe := queue.r.TxPipeline()
	//id delay 写入 zset
	pipe.ZAdd(ctx, queue.zSetKey, &redis.Z{Score: float64(absTime.UnixMilli()), Member: task.TaskID})
	//id arg 写入kv
	pipe.HSet(ctx, task.TaskID, "name", task.Name, "arg", task.Arg)
	_, err := pipe.Exec(ctx)
	if err != nil {
		return err
	}
	return nil
}

func (queue *TaskQueue) CancelTask(ctx context.Context, taskID string) error {
	return cancelTaskScript.Run(ctx, queue.r,
		[]string{queue.streamKey, taskID, queue.zSetKey}).Err()
}

func (queue *TaskQueue) StartConsumer() {
	queue.wg.Add(1)
	go queue.fetchTask()
}

func (queue *TaskQueue) Close() {
	close(queue.exit)
	queue.wg.Wait()
}

func (queue *TaskQueue) directlyAddTask(ctx context.Context, task *Task) error {
	return addStreamAndSetStreamIDScript.
		Run(ctx, queue.r, []string{queue.streamKey, task.TaskID}, task.Name, task.Arg).Err()
}

func (queue *TaskQueue) schedule() {
	defer queue.wg.Done()
loop:
	for {
		select {
		case <-time.Tick(queue.interval):
			now := time.Now().UTC()
			scheduleScript.Run(context.Background(), queue.r,
				[]string{},
				now.UnixMilli(),
				queue.batchCount,
			).Result()
		case <-queue.exit:
			break loop
		}
	}
}

func (queue *TaskQueue) fetchTask() {
	defer queue.wg.Done()
	ctx := context.Background()
loop:
	for {
		select {
		case <-time.Tick(queue.interval / 2):
			task, err := queue.readOneTask(ctx)
			if err != nil {
				continue loop
			}
			queue.do(task)
			queue.ack(task)
		case <-queue.exit:
			break loop
		}
	}
}

func (queue *TaskQueue) readOneTask(ctx context.Context) (*Task, error) {
	var (
		streams []redis.XStream
		err     error
	)
	for i := 0; i < 2; i++ {
		streams, err = queue.r.XReadGroup(ctx, &redis.XReadGroupArgs{
			Streams:  []string{queue.streamKey, ">"},
			Group:    queue.streamGroupName,
			Consumer: queue.streamConsumerName,
			Count:    1,
		}).Result()
		if err != nil {
			if strings.HasPrefix(err.Error(), "NOGROUP") {
				queue.createGroup()
			}
			queue.logger.Printf("XReadGroup: %v", err)
			continue
		}
		break
	}
	if err != nil {
		queue.logger.Printf("XReadGroup: %v", err)
		return nil, err
	}
	if len(streams[0].Messages) == 0 {
		return nil, nil
	}
	msg := streams[0].Messages[0]
	taskID, _ := msg.Values["task_id"].(string)
	m, err := queue.r.HGetAll(ctx, taskID).Result()
	if err != nil {
		return nil, err
	}
	task := &Task{
		TaskID:   taskID,
		StreamID: msg.ID,
		Name:     m["name"],
		Arg:      []byte(m["arg"]),
	}
	return task, err
}
func (queue *TaskQueue) createGroup() {
	queue.r.XGroupCreateMkStream(
		context.Background(),
		queue.streamKey,
		queue.streamGroupName,
		"0").Err()
}

func (queue *TaskQueue) do(task *Task) {
	v, ok := taskHandlerM.Load(task.Name)
	if !ok {
		queue.logger.Println("do: not found handler")
		return
	}
	handler, _ := v.(*TaskHandler)
	if handler.Recover != nil {
		defer handler.Recover()
	}
	handler.Handle(task.Arg)
}

func (queue *TaskQueue) ack(task *Task) {
	pipe := queue.r.Pipeline()
	//ack
	pipe.XAck(context.Background(), queue.streamKey, queue.streamGroupName, task.StreamID)
	pipe.Del(context.Background(), task.TaskID)
	_, err := pipe.Exec(context.Background())
	if err != nil {
		queue.logger.Println("ack(%+v): %v", task, err)
	}
}

func (queue *TaskQueue) pending() {
	defer queue.wg.Done()
loop:
	for {
		select {
		case <-queue.exit:
			break loop
		}
	}
}

func (queue *TaskQueue) cleanZombie() {
	defer queue.wg.Done()
loop:
	for {
		select {
		case <-queue.exit:
			break loop
		}
	}
}

func makeConsumerName() string {
	host, _ := os.Hostname()
	return fmt.Sprintf("%s_%d_%d", host, os.Getpid(), time.Now().UnixNano())
}
